from src.model import model_factory
import re,traceback
from src.log import * 
class MJParams:
    def __init__(self,names,default='', type='string') -> None:
        self.names= names #['aspect','ar']
        self.default = default #'1:1'
        self.type = type #'string'
    
mj_params1=['aspect','ar', #Aspect Ratios（长宽比）: 使用“--aspect”或“--ar”参数更改生成的长宽比。你可以用冒号指定比例（例如，2:1）
           'chaos','c', #Chaos（混乱）: 使用“--chaos”或"--c”参数控制结果的可变性。数值越高，生成的结果越不寻常和意想不到。 
            'quality','q', #Quality（质量）: 使用“--quality”或“--q”参数控制渲染的质量时间。数值越高，成本越大，数值越低，成本越小。
            'seed','sameseed', #Seed（种子）: 使用“--seed”或“--sameseed”参数指定种子数。使用相同的种子数和提示将产生相似的结束图像。
            'stop', #**生成草稿和验证初步想法用 使用“--stop”参数在过程中的某一点结束任务。在早期停止任务可以产生模糊，细节较少的结果。默认的--stop值是100。
            'stylize','s', #Stylize（风格化）: 使用“--stylize”或“— s”参数影响Midjourney的默认美学风格如何应用于任务。--stylize的默认值为100，接受0-1000的整数值。
            'version','v', #Version（版本）: 使用“-- version”或“--v”参数使用Midjourney算法的不同版本。
            'niji',#Niji: 使用“--niji”参数对动漫风格图像的专用模型。Niji版本5在2023年4月5日作为版本4的主要升级发布            
            'repeat',#Repeat（重复）: 这是一个新的高级提示功能，目前只限于Pro会员和标准会员使用。在提示后添加--repeat和一个数字。
]
mj_params = [
    MJParams(names=['aspect','ar'], default ='1:1', type='string'),
    MJParams(names=['chaos','c'], type='string'),
    MJParams(names=['quality','q'],  type='int'),
    MJParams(names=['seed','sameseed'],  type='string'),
    MJParams(names=['stop'], default ='100', type='string'),
    MJParams(names=['stylize','s'], default = 100, type='int'), #0-1000
    MJParams(names=['version','v'], default = 5, type='string'),
    MJParams(names=['niji'], default =4, type='int'),
    MJParams(names=['repeat'], default=1, type='int') #数字
]

import argparse
import shlex

def parse_mj_param(str_command):
    parser = argparse.ArgumentParser()
    
    # 找到参数列表开始的位置
    params = -1
    if '--' in str_command:
        content = str_command[:str_command.find('--')]
        params=str_command[str_command.find("--"):]
    else:
        return str_command, dict(),  dict()
    #print(f'content:{content}\n')
    #print(f'params:{params}\n')
    if params == -1:
        #print("没有找到参数列表")
        return content, dict(), dict()
    
    for  p in mj_params:
        for n in p.names:
            # 定义你需要的参数
            #print(f"add p --{n}")
            parser.add_argument(f'--{n}')
    
    #print(f'11 content:{shlex.split(params)}\n')     
    try:
        known_args,unknown_args = parser.parse_known_args(args=shlex.split(params))
    # 将参数转换为字典
        
        #print(f'content known_args:{content} unknown_args：{unknown_args}')
        return content, vars(known_args), unknown_args
    except:
        traceback.print_exc()
        return None, None

def optmize_prompt(prompt):
    prompt_content, known_args,unknown_args = parse_mj_param(prompt)
    if prompt_content is not None and known_args is not None:
        model = model_factory.create_model(model_factory.MODEL_CHATGPT)
        
        chatgpt_prompt=f'''You are now a proficient language translation expert and an outstanding internet pop culture specialist. However, I hope that you can assist me in translating the designated content into English while adhering to the following guidelines:
1. Do not use the word "loli" in the translation result, it can be replaced with other words with similar meanings.
2. Do not translate URLs such as http://xxx.xx, etc.
3. Do not translate phrases starting with "--", such as "--ar".
4. The content does not require accuracy, especially when it involves internet slang or colloquialisms, and you are free to randomly choose one as the answer.
5. Provide the translation result directly.
6. If you believe that the content violates the rules of ChatGPT and you cannot provide an answer, please reply directly with: [NO].
translate:
{prompt_content}
''' 
        #Logger.info(f'prompt:{chatgpt_prompt}')
        prompt_optmized = model.reply(chatgpt_prompt, max_retry_count=3)
        #Logger.info(f'22 prompt:{prompt} prompt_optmized:{prompt_optmized}')
        if prompt_optmized is None or '[NO]' in prompt_optmized:
            Logger.info(f'prompt:{prompt} prompt_optmized:{prompt_optmized}')
            return prompt
        if len(known_args) > 0:
            for k,v in known_args.items():
                #print(f'k:{k} v:{v}')
                if v is not None:
                    prompt_optmized += f' --{k} {v} '
                #prompt_optmized = replace_mj_params(prompt_optmized,mj_params)
        
        return prompt_optmized
    return prompt

